// GPEncryptDlg.cpp : implementation file
//

#include "stdafx.h"
#include "GPEncrypt.h"
#include "GPEncryptDlg.h"
#include "ExchangeData.h"
#include "FileIO.h"
#include "GPGMEExt.h"

#include "ProgressDlg.h"

#ifdef _DEBUG
#define DEBUG_DATA
#endif

#ifdef _DEBUG
#define new DEBUG_NEW
#endif

extern "C"
{
	void progress_cb (
		void *hook,
		const char *what,
		int type,
		int current,
		int total);
};

void progress_cb (
	void *hook,
	const char *what,
	int type,
	int current,
	int total)
{
	int total_length = *(int*)hook;
	int pos = 0;

	//Make sure 100 is finish.
	if (total_length == current)
	{
		pos = 100;
	}
	else
	{
		pos = (int)(current * 100.0f /
			total_length);
		pos = (pos == 100)?99:pos;
	}

	*pexdata_percent_value = pos;
}

UINT gpgme_encrypt_proc(LPVOID* pParam)
{
	CGPEncryptDlg* dlg = (CGPEncryptDlg*)pParam;
	//File handle
	FILE_IO_META* hinfo_in  = NULL;
	FILE_IO_META* hinfo_out = NULL;

	//Data handle
	gpgme_data_t hdata_in	= NULL;
	gpgme_data_t hdata_out	= NULL;

	//Encrypt op result
	gpgme_error_t err;

	//Recp keys
	gpgme_key_t* keys = NULL;

	CString stringTmp;

	//////////////////////////////////////////////////////////
	//Open input file
#ifdef DEBUG_DATA
	stringTmp = "../__data__.bin";
#else
	stringTmp = __argv[1];
#endif

	if (gpgme_data_new_for_read(
		&hdata_in,
		stringTmp.GetBuffer(),
		&hinfo_in) !=
		GPG_ERR_NO_ERROR)
	{
		AfxMessageBox("Can't open input file!");
		return -1;
	}

	//Open output file
#ifdef DEBUG_DATA
	stringTmp = "../__data__.bin.gpg";
#else
	stringTmp.Format("%s.gpg",__argv[1]);
#endif

	if (gpgme_data_new_for_write(
		&hdata_out,
		stringTmp.GetBuffer(),
		&hinfo_out) !=
		GPG_ERR_NO_ERROR)
	{
		AfxMessageBox("Can't open output file!");
		gpgme_data_release(hdata_in);
		return -1;
	}

	//////////////////////////////////////////////////////////
	//Set callback function
	gpgme_set_progress_cb (
		dlg->ctx, 
		progress_cb,
		&(hinfo_in->nLowSize));

	//////////////////////////////////////////////////////////
	//Get keys
	int count = dlg->m_list_used_keys.GetItemCount();
	keys = (gpgme_key_t*)malloc(sizeof(gpgme_key_t) * (count + 1));

	memset(keys,0,sizeof(gpgme_key_t) * (count + 1));

	for (int i = 0;i < count;i++)
	{
		keys[i] = 
			(gpgme_key_t)dlg->m_list_used_keys.GetItemData(i);
	}

	//////////////////////////////////////////////////////////
	//Encrypt process.
	err = gpgme_op_encrypt(
		dlg->ctx, 
		keys,
		GPGME_ENCRYPT_ALWAYS_TRUST,
		hdata_in,
		hdata_out);

	if (err != GPG_ERR_NO_ERROR)
	{
		stringTmp.Format("Error %s:%s\n",
			gpgme_strsource (err),
			gpgme_strerror (err));
		AfxMessageBox(stringTmp);
	}

	//////////////////////////////////////////////////////////
	//Clean data
	gpgme_data_release (hdata_in);
	gpgme_data_release (hdata_out);
	free(keys);

	//Exit program
	dlg->SendMessage(WM_CLOSE);

	return 0;
}

// CGPEncryptDlg dialog
CGPEncryptDlg::CGPEncryptDlg(CWnd* pParent /*=NULL*/)
	: CDialog(CGPEncryptDlg::IDD, pParent)
{
	m_hIcon = AfxGetApp()->LoadIcon(IDR_MAINFRAME);
	this->ctx = NULL;
	this->hTerminalProcess = CreateEvent(NULL,FALSE,FALSE,NULL);
}


CGPEncryptDlg::~CGPEncryptDlg()
{
	map<CString,gpgme_key_t>::iterator p;

	for (p = m_keys.begin();p != m_keys.end();p++)
	{
		gpgme_key_unref(p->second);
	}

	if (ctx != NULL)
	{
		gpgme_release(ctx);
		ctx = NULL;
	}

	CloseHandle(this->hTerminalProcess);
	this->hTerminalProcess = INVALID_HANDLE_VALUE;
}


void CGPEncryptDlg::DoDataExchange(CDataExchange* pDX)
{
	CDialog::DoDataExchange(pDX);
	DDX_Control(pDX, IDC_LIST_ALL_VAILD_KEYS, m_list_all_keys);
	DDX_Control(pDX, IDC_LIST_USED_KEYS, m_list_used_keys);
	DDX_Control(pDX, IDC_CHECK_TEXT_OUTPUT, m_check_text_output);
}

BEGIN_MESSAGE_MAP(CGPEncryptDlg, CDialog)
	ON_WM_PAINT()
	ON_WM_QUERYDRAGICON()
	//}}AFX_MSG_MAP
	ON_NOTIFY(NM_DBLCLK, IDC_LIST_ALL_VAILD_KEYS, &CGPEncryptDlg::OnNMDblclkListAllVaildKeys)
	ON_NOTIFY(NM_DBLCLK, IDC_LIST_USED_KEYS, &CGPEncryptDlg::OnNMDblclkListUsedKeys)
	ON_BN_CLICKED(IDC_BUTTON_OK, &CGPEncryptDlg::OnBnClickedButtonOk)
	ON_BN_CLICKED(IDC_CHECK_TEXT_OUTPUT, &CGPEncryptDlg::OnBnClickedCheckTextOutput)
	ON_MESSAGE(WM_USER_PROGRESS_DLG_CANCEL, OnProgressDlgCancel)
	ON_BN_CLICKED(IDC_BUTTON_CANCEL, &CGPEncryptDlg::OnBnClickedButtonCancel)
END_MESSAGE_MAP()


// CGPEncryptDlg message handlers

BOOL CGPEncryptDlg::OnInitDialog()
{
	CDialog::OnInitDialog();

	// Set the icon for this dialog.  The framework does this automatically
	//  when the application's main window is not a dialog
	SetIcon(m_hIcon, TRUE);			// Set big icon
	SetIcon(m_hIcon, FALSE);		// Set small icon

	// TODO: Add extra initialization here
	CRect rect;
	m_list_all_keys.SetExtendedStyle(
		m_list_all_keys.GetExtendedStyle()|
		LVS_EX_FULLROWSELECT);

	m_list_used_keys.SetExtendedStyle(
		m_list_used_keys.GetExtendedStyle()|
		LVS_EX_FULLROWSELECT);

	m_list_all_keys.GetClientRect(&rect);
	m_list_all_keys.InsertColumn(0,
		"User list (Double click user will add to recipients list.)", 0, int(rect.Width() * 0.65));
	m_list_all_keys.InsertColumn(1,
		"Valid", 0, int(rect.Width() * 0.15));
	m_list_all_keys.InsertColumn(2,
		"Size", 0, int(rect.Width() * 0.15));

	m_list_used_keys.GetClientRect(&rect);
	m_list_used_keys.InsertColumn(0,
		"Recipients", 0, int(rect.Width() * 0.65));
	m_list_used_keys.InsertColumn(1,
		"Valid", 0, int(rect.Width() * 0.15));
	m_list_used_keys.InsertColumn(2,
		"Size", 0, int(rect.Width() * 0.15));

	//Set icons
	m_key_icon.Create(22,14,ILC_COLOR8,0,0);
	m_key_icon.SetBkColor(RGB(255,255,255));
	m_key_icon.Add(
		LoadIcon(AfxGetResourceHandle(),
			MAKEINTRESOURCE(IDI_ICON_KEYPAIR)));
	m_key_icon.Add(
		LoadIcon(AfxGetResourceHandle(),
			MAKEINTRESOURCE(IDI_ICON_PUBKEY)));
	m_key_icon.Add(
		LoadIcon(AfxGetResourceHandle(),
		MAKEINTRESOURCE(IDI_ICON_EXP_KEYPAIR)));
	m_key_icon.Add(
		LoadIcon(AfxGetResourceHandle(),
		MAKEINTRESOURCE(IDI_ICON_EXP_PUBKEY)));
	m_key_icon.Add(
		LoadIcon(AfxGetResourceHandle(),
		MAKEINTRESOURCE(IDI_ICON_DISABLE_KEYPAIR)));
	m_key_icon.Add(
		LoadIcon(AfxGetResourceHandle(),
		MAKEINTRESOURCE(IDI_ICON_DISABLE_PUBKEY)));

	m_list_all_keys.SetImageList(&m_key_icon,LVSIL_SMALL);
	m_list_used_keys.SetImageList(&m_key_icon,LVSIL_SMALL);

	/* init gpgme */
	InitGPGMESystem (GPGME_PROTOCOL_OpenPGP);

	/* create our own context */
	gpgme_error_t err = gpgme_new (&ctx);
	if (err != GPG_ERR_NO_ERROR)
	{
		OutputDebugString("error of gpgme_new\n");
	}
	else
	{
		/* set protocol to use in our context */
		err = gpgme_set_protocol(ctx, GPGME_PROTOCOL_OpenPGP);
		if (err != GPG_ERR_NO_ERROR)
		{
			OutputDebugString("error of gpgme_set_protocol\n");
		}
		else
		{
			/* do ascii armor data, so output is readable in console */
			gpgme_set_armor(ctx, 1);
			m_check_text_output.SetCheck(TRUE);

			/* get gpg user's keys and display in listctrl */
			GetPrivateKeysList();
			GetKeysInfoToListctrl();
		}
	}

	return TRUE;  // return TRUE  unless you set the focus to a control
}

// If you add a minimize button to your dialog, you will need the code below
//  to draw the icon.  For MFC applications using the document/view model,
//  this is automatically done for you by the framework.

void CGPEncryptDlg::OnPaint()
{
	if (IsIconic())
	{
		CPaintDC dc(this); // device context for painting

		SendMessage(WM_ICONERASEBKGND, reinterpret_cast<WPARAM>(dc.GetSafeHdc()), 0);

		// Center icon in client rectangle
		int cxIcon = GetSystemMetrics(SM_CXICON);
		int cyIcon = GetSystemMetrics(SM_CYICON);
		CRect rect;
		GetClientRect(&rect);
		int x = (rect.Width() - cxIcon + 1) / 2;
		int y = (rect.Height() - cyIcon + 1) / 2;

		// Draw the icon
		dc.DrawIcon(x, y, m_hIcon);
	}
	else
	{
		CDialog::OnPaint();
	}
}

// The system calls this function to obtain the cursor to display while the user drags
//  the minimized window.
HCURSOR CGPEncryptDlg::OnQueryDragIcon()
{
	return static_cast<HCURSOR>(m_hIcon);
}

void CGPEncryptDlg::OnNMDblclkListAllVaildKeys(NMHDR *pNMHDR, LRESULT *pResult)
{
	// TODO: Add your control notification handler code here
	*pResult = 0;
	MoveListItem(m_list_used_keys,m_list_all_keys);
}

void CGPEncryptDlg::OnNMDblclkListUsedKeys(NMHDR *pNMHDR, LRESULT *pResult)
{
	// TODO: Add your control notification handler code here
	*pResult = 0;
	MoveListItem(m_list_all_keys, m_list_used_keys);
}

void CGPEncryptDlg::OnBnClickedButtonOk()
{
	// TODO: Add your control notification handler code here

	//Create and init temp variable
	CProgressDlg dlg("Encrypt",this);

	//Prepare sync values
	*pexdata_percent_value = 0;

#ifdef DEBUG_DATA
#else
	//Check filename
	if (__argc != 2)
	{
		AfxMessageBox("Invalid file name!");
		return;
	}
#endif

	if (m_list_used_keys.GetItemCount() == 0)
	{
		AfxMessageBox("You must choose one key at least!");
		return;
	}

	//Create worker thread
	AfxBeginThread((AFX_THREADPROC)gpgme_encrypt_proc, this);

	//Pending...
	this->ShowWindow(FALSE);
	dlg.DoModal();
}

void CGPEncryptDlg::OnBnClickedCheckTextOutput()
{
	// TODO: Add your control notification handler code here
	if (m_check_text_output.GetCheck() == BST_CHECKED)
	{
		//Output is binary, for save size.
		gpgme_set_armor(ctx, 1);
	}
	else
	{
		//Output is text, human readable
		gpgme_set_armor(ctx, 0);
	}
}

void CGPEncryptDlg::OnBnClickedButtonCancel()
{
	// TODO: Add your control notification handler code here
	this->SendMessage(WM_CLOSE);
}

LRESULT CGPEncryptDlg::OnProgressDlgCancel(
	WPARAM wparam,
	LPARAM lparam)
{
	//Kill gpg.exe
	PID_LIST* pid_list = FindMyChildProcessIDs();
	if (pid_list != NULL)
	{
		for (PID_LIST::iterator p = pid_list->begin();
			p != pid_list->end();
			p++)
		{
			if (strcmp(p->szExeFile,"gpg.exe") == 0)
			{
				OutputDebugString("Kill gpg.exe\n");
				KillProcessByID(p->th32ProcessID);
			}
		}

		delete pid_list;
	}

	return 0;
}

void CGPEncryptDlg::InitGPGMESystem(gpgme_protocol_t proto)
{
	gpgme_error_t err;

	gpgme_check_version (NULL);
	setlocale (LC_ALL, "");
	gpgme_set_locale (NULL, LC_CTYPE, setlocale (LC_CTYPE, NULL));

	err = gpgme_engine_check_version (proto);
}

void CGPEncryptDlg::GetKeysInfoToListctrl()
{
	gpgme_error_t			err;
	gpgme_key_t             key;

	/* get all user keys. */
	err = gpgme_op_keylist_start(ctx, NULL, FALSE);
	while (!(err = gpgme_op_keylist_next(ctx, &key)))
	{
		int index = 0;
		CString stringTmp;

		if (key->uids == NULL)
		{
			break;
		} else if ((key->uids->name == NULL) ||
			(key->uids->email == NULL) ||
			(key->uids->uid == NULL))
		{
			continue;
		}
		else
		{
			index = m_list_all_keys.GetItemCount();
		}

		//Insert into key listctrl
		stringTmp.Empty();
		GenerateUserInfoString(key, stringTmp);
		GPGME_KEY_TYPE key_type = DetermineKeyType(key);
		m_list_all_keys.InsertItem(index,stringTmp,key_type);

		//Bind data
		m_keys[key->uids->uid] = key;
		m_list_all_keys.SetItemData(index, (DWORD_PTR)key);

		//set key status
		stringTmp.Empty();
		GenerateKeyStatus(key, stringTmp);
		m_list_all_keys.SetItemText(index,1,stringTmp);

		//set subkeys length.
		stringTmp.Empty();
		GenerateSubkeysInfoString(key->subkeys, stringTmp);
		m_list_all_keys.SetItemText(index,2,stringTmp);
	}
}

void CGPEncryptDlg::GetPrivateKeysList()
{
	/* get secret key */
	gpgme_error_t			err;
	gpgme_key_t             key;
	err = gpgme_op_keylist_start(ctx, NULL, TRUE);

	while (!(err = gpgme_op_keylist_next(ctx, &key)))
	{
		if (key->uids == NULL)
		{
			break;
		}
		else
		{
			m_keys_is_secret[key->uids->uid] = true;
		}

		gpgme_key_unref(key);
	}
}

void CGPEncryptDlg::GenerateUserInfoString(gpgme_key_t key, CString &stringTmp)
{
	if ((key->uids->comment != NULL) && (strlen(key->uids->comment) > 0))
	{
		stringTmp.Format("%s <%s> (%s)", 
			key->uids->name,
			key->uids->email,
			key->uids->comment);
	}
	else
	{
		stringTmp.Format("%s <%s>", 
			key->uids->name,
			key->uids->email);
	}
}

void CGPEncryptDlg::GenerateKeyStatus(gpgme_key_t key, CString& stringTmp)
{
	if (key->disabled)
	{
		stringTmp = "disable";
	}
	else
	{
		switch (key->uids->validity)
		{
		case GPGME_VALIDITY_UNKNOWN:
			stringTmp = "unknown";
			break;
		case GPGME_VALIDITY_UNDEFINED:
			stringTmp = "undefined";
			break;
		case GPGME_VALIDITY_NEVER:
			stringTmp = "never";
			break;
		case GPGME_VALIDITY_MARGINAL:
			stringTmp = "marginal";
			break;
		case GPGME_VALIDITY_FULL:
			stringTmp = "full";
			break;
		case GPGME_VALIDITY_ULTIMATE:
			stringTmp = "ultimate";
			break;
		default:
			stringTmp = "error";
			break;
		}
	}
}

void CGPEncryptDlg::GenerateSubkeysInfoString(gpgme_subkey_t subkeys, CString& subkeys_info)
{
	CString stringTmp;	
	gpgme_subkey_t subkey_next;

	subkey_next = subkeys;

	while (subkey_next != NULL) {
		if (subkey_next->fpr != NULL) {
			stringTmp.Format("%d",subkey_next->length);
		}

		if (subkeys_info.GetLength() > 0)
		{
			subkeys_info += "/" + stringTmp;
		}
		else
		{
			subkeys_info = stringTmp;
		}

		subkey_next = subkey_next->next;
	}
}

void CGPEncryptDlg::MoveListItem(CListCtrl& list_dst, CListCtrl& list_src)
{
	int indexOfSrc = 0;
	int indexOfDst = 0;

	if (&list_src == &list_dst)
	{
		return;
	}

	if (list_src.GetItemCount() == 0)
	{
		return;
	}

	//Find item
	for (int i = 0;i < list_src.GetItemCount();i++)
	{
		if (list_src.GetItemState(i,LVIS_SELECTED) == LVIS_SELECTED)
		{
			indexOfSrc = i;
			break;
		}
	}

	//Copy item text
	CString stringTmp;
	gpgme_key_t key;

	indexOfDst = list_dst.GetItemCount();

	//Get src text and data
	stringTmp = list_src.GetItemText(indexOfSrc,0);
	key = (gpgme_key_t)list_src.GetItemData(indexOfSrc);

	if ((key == NULL) || 
		(key->uids == NULL) || 
		(key->uids->uid == NULL))
	{
		OutputDebugString("list data : key is invalid\n");
		return;
	}

	GPGME_KEY_TYPE key_type = DetermineKeyType(key);

	list_dst.InsertItem(indexOfDst, stringTmp,key_type);
	list_dst.SetItemData(indexOfDst,(DWORD_PTR)key);

	stringTmp = list_src.GetItemText(indexOfSrc,1);
	list_dst.SetItemText(indexOfDst,1,stringTmp);

	stringTmp = list_src.GetItemText(indexOfSrc,2);
	list_dst.SetItemText(indexOfDst,2,stringTmp);

	//Delete from all keys.
	list_src.DeleteItem(indexOfSrc);
}

GPGME_KEY_TYPE CGPEncryptDlg::DetermineKeyType( gpgme_key_t key )
{
	GPGME_KEY_TYPE key_type;

	if (m_keys_is_secret.find(key->uids->uid) != 
		m_keys_is_secret.end())
	{
		//have secret key
		if (key->expired)
		{
			key_type = ICON_EXP_KEYPAIR;
		}
		else if (key->disabled)
		{
			key_type = ICON_DISABLE_KEYPAIR;
		}
		else
		{
			key_type = ICON_KEYPAIR;
		}
	}
	else
	{
		//public key only
		if (key->expired)
		{
			key_type = ICON_EXP_PUBKEY;
		}
		else if (key->disabled)
		{
			key_type = ICON_DISABLE_PUBKEY;
		}
		else
		{
			key_type = ICON_PUBKEY;
		}
	}

	return key_type;
}

